 /*
	This file is part of PropExtractor.

	PropExtractor is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	PropExtractor is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with PropExtractor. If not, see <http://www.gnu.org/licenses/>.
*/

/*
	The decryption for QNT images was originally written by
	Masaki Chikama in 2001. I have adapted it to C++ so it
	can continue to be used and adapted.

	I thank Mr. Chikama for making his code freely available.

	-MrHands
*/

#include "PEIncludes.h"

#include "PEAliceSoftImageQnt.h"

#include "PETools.h"
#include "PEAliceSoftArchiveRed.h"

namespace pe
{

	// ================================
	// Image
	// ================================

	int AliceImageQntReadHeader(AliceImageQntHeader* a_Header, PEStream* a_Stream)
	{
		PEStreamRead(a_Stream, a_Header, sizeof(AliceImageQntHeader), 1);

		size_t header_size = sizeof(AliceImageQntHeader);

		if (a_Header->header[0] != 'Q' ||
			a_Header->header[1] != 'N' ||
			a_Header->header[2] != 'T')
		{
			return -1;
		}

		return PE_OK;
	}

	int AliceImageQntReadData(AliceImageQnt* a_Target, AliceImageQntHeader* a_Header, PEStream* a_Stream)
	{
		a_Target->header_v0 = NULL;
		a_Target->header_v2 = NULL;

		if (a_Header->rsv0 == 0)
		{
			a_Target->header_v0 = new AliceImageQntHeaderV0;
			PEStreamRead(a_Stream, a_Target->header_v0, sizeof(AliceImageQntHeaderV0), 1);

			a_Target->width = a_Target->header_v0->width;
			a_Target->height = a_Target->header_v0->height;
			a_Target->bpp = (byte)(a_Target->header_v0->bpp >> 3);

			a_Target->pixel_size = a_Target->header_v0->pixel_size;
			a_Target->alpha_size = 0;
		}
		else if (a_Header->rsv0 == 1)
		{
			a_Target->header_v1 = new AliceImageQntHeaderV1;
			PEStreamRead(a_Stream, a_Target->header_v1, sizeof(AliceImageQntHeaderV1), 1);

			a_Target->width = a_Target->header_v1->width;
			a_Target->height = a_Target->header_v1->height;
			a_Target->bpp = (byte)(a_Target->header_v1->bpp >> 3);

			a_Target->pixel_size = a_Target->header_v1->pixel_size;
			a_Target->alpha_size = 0;
		}
		else if (a_Header->rsv0 == 2)
		{
			a_Target->header_v2 = new AliceImageQntHeaderV2;
			PEStreamRead(a_Stream, a_Target->header_v2, sizeof(AliceImageQntHeaderV2), 1);

			a_Target->width = a_Target->header_v2->width;
			a_Target->height = a_Target->header_v2->height;
			a_Target->bpp = (byte)(a_Target->header_v2->bpp >> 3);

			a_Target->pixel_size = a_Target->header_v2->pixel_size;
			a_Target->alpha_size = a_Target->header_v2->alpha_size;
		}
		else
		{
			PE_PRINT("ERROR! Unhandled rsv0: %i", a_Header->rsv0);
			return -1;
		}

		a_Target->header = a_Header;
		
		a_Target->alpha = NULL;

		if (a_Target->bpp == 0) { return -1; }

		PE_DEBUG("Reading %s...", a_Stream->path);

		z_stream zlib_stream;
		zlib_stream.zalloc = Z_NULL;
		zlib_stream.zfree = Z_NULL;
		zlib_stream.opaque = Z_NULL;
		zlib_stream.avail_in = 0;
		zlib_stream.next_in = Z_NULL;

		inflateInit(&zlib_stream);

		byte* dst;

		unsigned int total_size = a_Target->width * a_Target->height * a_Target->bpp + (5 * 1024);
		a_Target->pixels = new byte[total_size];
		memset(a_Target->pixels, 0, total_size);
		dst = a_Target->pixels;

		unsigned int total_read = 0;
		unsigned int total_compressed = a_Target->pixel_size + (5 * 1024);

		byte* src_buffer = new byte[a_Target->pixel_size + a_Target->alpha_size];
		memset(src_buffer, 0, a_Target->pixel_size + a_Target->alpha_size);
		PEStreamRead(a_Stream, src_buffer, 1, a_Target->pixel_size + a_Target->alpha_size);

		uncompress(
			a_Target->pixels, 
			(unsigned long*)&total_size, 
			src_buffer, 
			total_compressed
		);

		PEStreamClose(a_Stream);

		PE_DEBUG("Alpha size: %i", a_Target->alpha_size);
		if (a_Target->alpha_size > 0)
		{
			a_Target->alpha_data_size = a_Target->width * a_Target->height;
			a_Target->alpha = new byte[a_Target->alpha_data_size + (10 * 1024)];
			uncompress(
				a_Target->alpha, 
				(unsigned long*)&a_Target->alpha_data_size, 
				src_buffer + a_Target->pixel_size,
				a_Target->alpha_size
			);
		}
		else
		{
			a_Target->alpha = NULL;
		}

		delete src_buffer;

		PE_DEBUG("Image dimensions: %i x %i", a_Target->width, a_Target->height);

		PE_DEBUG("Image size: %i", total_read);

		return 0;
	}

	int AliceImageQntToRaw(PEImageRaw* a_Target, AliceImageQnt* a_Image)
	{
		PEImageRawCreate(a_Target, a_Image->width, a_Image->height);

		// reverse obfuscation

		byte* src_r = a_Image->pixels;
		byte* src_g = a_Image->pixels + (a_Image->width * a_Image->height * 1);
		byte* src_b = a_Image->pixels + (a_Image->width * a_Image->height * 2);

		byte* temp_color = new byte[a_Image->width * a_Image->height * 4];

		byte* dst = temp_color;

		byte* read_r;
		byte* read_g;
		byte* read_b;

		byte* dst_ul;
		byte* dst_ll;
		byte* dst_ur;
		byte* dst_lr;

		unsigned int x, y;

		unsigned int pixel = 0;

		for (int i = 2; i >= 0; i--)
		{
			byte* dst = temp_color;
			byte* read = &a_Image->pixels[pixel];

			for (y = 0; y < a_Image->height - 1; y += 2)
			{
				read_r = src_r + (y * a_Image->width);
				read_g = src_g + (y * a_Image->width);
				read_b = src_b + (y * a_Image->width);

				dst_ul = dst;
				dst_ll = dst + (a_Image->width * 3);
				dst_ur = dst + 3;
				dst_lr = dst + ((a_Image->width + 1) * 3);

				for (x = 0; x < a_Image->width - 1; x += 2)
				{						
					dst_ul[2 - i] = read[0];
					dst_ll[2 - i] = read[1];
					dst_ur[2 - i] = read[2];
					dst_lr[2 - i] = read[3];

					dst_ul += 6;
					dst_ll += 6;
					dst_ur += 6;
					dst_lr += 6;

					read += 4;

					pixel += 4;
				}
				if (x != a_Image->width)
				{
					temp_color[( y      * a_Image->width + x) * 3 + i] = read[0];
					temp_color[((y + 1) * a_Image->width + x) * 3 + i] = read[1];

					read += 4;

					pixel += 4;
				}

				dst += a_Image->width * 6;
			}

			if (y != a_Image->height)
			{
				dst = temp_color;

				for (x = 0; x < a_Image->width; x += 2)
				{
					if (x + 1 < a_Image->width)
					{
						temp_color[(y * a_Image->width + x    ) * 3 + i] = read[0];
						temp_color[(y * a_Image->width + x + 1) * 3 + i] = read[2];

						read += 4;

						pixel += 4;
					}
					else
					{
						temp_color[(y * a_Image->width + x    ) * 3 + i] = read[0];

						read += 2;

						pixel += 2;
					}
				}
			}
		}

		for (unsigned int x = 1; x < a_Image->width; x++)
		{
			temp_color[x * 3 + 0] = temp_color[(x - 1) * 3 + 0] - temp_color[x * 3 + 0];
			temp_color[x * 3 + 1] = temp_color[(x - 1) * 3 + 1] - temp_color[x * 3 + 1];
			temp_color[x * 3 + 2] = temp_color[(x - 1) * 3 + 2] - temp_color[x * 3 + 2];
		}

		for (unsigned int y = 1; y < a_Image->height; y++)
		{
			temp_color[(y * a_Image->width) * 3 + 0] = temp_color[((y - 1) * a_Image->width) * 3 + 0] - temp_color[(y * a_Image->width) * 3 + 0];
			temp_color[(y * a_Image->width) * 3 + 1] = temp_color[((y - 1) * a_Image->width) * 3 + 1] - temp_color[(y * a_Image->width) * 3 + 1];
			temp_color[(y * a_Image->width) * 3 + 2] = temp_color[((y - 1) * a_Image->width) * 3 + 2] - temp_color[(y * a_Image->width) * 3 + 2];

			for (unsigned int x = 1; x < a_Image->width; x++)
			{
				byte px, py;

				py = temp_color[((y - 1) * a_Image->width + (x    )) * 3 + 0];
				px = temp_color[((y    ) * a_Image->width + (x - 1)) * 3 + 0];
				temp_color[(y * a_Image->width + x) * 3 + 0] = ((py + px) >> 1) - temp_color[(y * a_Image->width + x) * 3 + 0];

				py = temp_color[((y - 1) * a_Image->width + (x    )) * 3 + 1];
				px = temp_color[((y    ) * a_Image->width + (x - 1)) * 3 + 1];
				temp_color[(y * a_Image->width + x) * 3 + 1] = ((py + px) >> 1) - temp_color[(y * a_Image->width + x) * 3 + 1];

				py = temp_color[((y - 1) * a_Image->width + (x    )) * 3 + 2];
				px = temp_color[((y    ) * a_Image->width + (x - 1)) * 3 + 2];
				temp_color[(y * a_Image->width + x) * 3 + 2] = ((py + px) >> 1) - temp_color[(y * a_Image->width + x) * 3 + 2];
			}
		}


		// alpha

		byte* temp_alpha = new byte[a_Image->width * a_Image->height * 2];
		memset(temp_alpha, 0xFF, a_Image->width * a_Image->height * 2);
		//a_Image->alpha_data_size = a_Image->width * a_Image->height;

		if (a_Image->alpha_size > 0)
		{
			byte* dst_alpha = temp_alpha;

			byte* src_alpha = a_Image->alpha;

			unsigned int curr_alpha = 1;
			dst_alpha[0] = src_alpha[0];
			for (unsigned int x = 1; x < a_Image->width; x++)
			{
				dst_alpha[x] = dst_alpha[x - 1] - src_alpha[curr_alpha];
				curr_alpha++;
			}

			for (unsigned int y = 1; y < a_Image->height; y++)
			{
				dst_alpha[y * a_Image->width] = dst_alpha[(y - 1) * a_Image->width] - src_alpha[curr_alpha];
				curr_alpha++;

				for (unsigned int x = 1; x < a_Image->width; x++)
				{
					byte px = dst_alpha[ (y      * a_Image->width) + x - 1];
					byte py = dst_alpha[((y - 1) * a_Image->width) + x    ];
					dst_alpha[y * a_Image->width + x] = ((px + py) >> 1) - src_alpha[curr_alpha];
					curr_alpha++;
				}
			}
		}

		//a_Target->data = new byte[a_Image->width * a_Image->height * 4];

		byte* src_color = temp_color;
		byte* src_alpha = temp_alpha;
		byte* dst_pixel = a_Target->data;
		for (unsigned int i = 0; i < a_Image->width * a_Image->height; i++)
		{
			dst_pixel[0] = src_color[0];
			dst_pixel[1] = src_color[1];
			dst_pixel[2] = src_color[2];
			dst_pixel[3] = src_alpha[0];

			dst_pixel += 4;

			src_color += 3;
			src_alpha++;
		}

		delete temp_color;
		delete temp_alpha;

		return 0;
	}

	int AliceImageQntDeleteHeader(AliceImageQntHeader* a_Header)
	{
		return PE_OK;
	}

	int AliceImageQntDeleteData(AliceImageQnt* a_Image)
	{
		if (a_Image->alpha_size > 0)
		{
			delete a_Image->alpha;
		}
		delete a_Image->pixels;
		a_Image->data_size = 0;

		a_Image->width = a_Image->height = 0;
		a_Image->bpp = 0;
		a_Image->alpha_size = 0;

		return PE_OK;
	}

	int AliceImageQntReadPathAndDecrypt(PEImageRaw* a_Target, const wchar_t* a_Path)
	{
		PEStream read;
		PEStreamOpen(&read, a_Path);

		return AliceImageQntReadAndDecrypt(a_Target, &read);
	}

	int AliceImageQntReadAndDecrypt(PEImageRaw* a_Target, PEStream* a_Stream)
	{
		AliceImageQntHeader header;
		if (AliceImageQntReadHeader(&header, a_Stream) != 0) 
		{
			PE_PRINT("ERROR! Not an AliceSoft QNT compressed image.");
			AliceImageQntDeleteHeader(&header);
			PEStreamClose(a_Stream);
			return -1;
		}

		AliceImageQnt image;
		if (AliceImageQntReadData(&image, &header, a_Stream) != PE_OK)
		{
			PE_PRINT("ERROR! Not an AliceSoft QNT compressed image.");
			AliceImageQntDeleteHeader(&header);
			AliceImageQntDeleteData(&image);
			PEStreamClose(a_Stream);
			return -1;
		}

		if (AliceImageQntToRaw(a_Target, &image) != PE_OK)
		{
			PE_PRINT("ERROR! Not an AliceSoft QNT compressed image.");
			AliceImageQntDeleteHeader(&header);
			AliceImageQntDeleteData(&image);
			PEStreamClose(a_Stream);
			return -1;
		}

		AliceImageQntDeleteHeader(&header);
		AliceImageQntDeleteData(&image);
		PEStreamClose(a_Stream);

		return PE_OK;
	}

	int AliceImageQntDecryptAndSave(const wchar_t* a_Path)
	{
		PEStream read;
		PEStreamOpen(&read, a_Path);

		AliceImageQntHeader header;
		if (AliceImageQntReadHeader(&header, &read) != 0) 
		{
			PE_PRINT("ERROR! Not an AliceSoft compressed image.");
			return -1;
		}

		AliceImageQnt image;
		AliceImageQntReadData(&image, &header, &read);

		PEImageRaw raw;
		AliceImageQntToRaw(&raw, &image);

		PEImageRawSaveToFile(&raw, read.path);

		PEImageRawDelete(&raw);
		AliceImageQntDeleteData(&image);
		AliceImageQntDeleteHeader(&header);

		PEStreamClose(&read);

		return PE_OK;
	}


	ImageAliceSoftQnt::ImageAliceSoftQnt() : PEImage()
	{

	}

	ImageAliceSoftQnt::~ImageAliceSoftQnt()
	{

	}

	int ImageAliceSoftQnt::LoadHeader(PEStream* a_Stream)
	{
		m_Stream = a_Stream;
		if (AliceImageQntReadHeader(&m_Header, a_Stream) == PE_OK)
		{
			return PE_OK;
		}

		return -1;
	}

	int ImageAliceSoftQnt::LoadBody()
	{
		AliceImageQnt qnt_data;
		AliceImageQntReadData(&qnt_data, &m_Header, m_Stream);

		m_Data = new PEImageRaw;
		AliceImageQntToRaw(m_Data, &qnt_data);

		AliceImageQntDeleteData(&qnt_data);
		AliceImageQntDeleteHeader(&m_Header);

		return PE_OK;
	}

}; // namespace pe